# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/5 14:59
# @Function: 同过土壤保持的实物量来计算价值量
"""
    通过土壤保持的实物量来计算价值量。
输入：
    Qsr：土壤保持的实物量：某文件路径       数据类型：tif
    泥沙淤积系数：0.24        数据类型：常量
    土壤容重：1.27        数据类型：常量
    水库单位清淤工程费用：60        数据类型：常量
    sr_TN: 污染物氮纯含量：某文件路径     数据类型：tif
    sr_TP: 污染物磷纯含量：某文件路径     数据类型：tif
    污染物氮单位处理成本：3000        数据类型：常量
    污染物磷单位处理成本：3500        数据类型：常量
输出：
    Vsr：土壤保持价值      数据类型：tif        计算方法：Vsr = Vsd + Vdpd
    Vsd: 减少泥沙淤积价值       数据类型：tif        计算方法：Vsd = 泥沙淤积系数 * Qsr / 土壤容重 * 水库单位清淤工程费用
    Vdpd: 减少面源污染价值      数据类型：tif        计算方法：Vdpd = Qsr *（ sr_TN * 污染物氮单位处理成本 + sr_TP * 污染物磷单位处理成本）
"""
# 数据的前期准备
# 根据shp裁剪保证Qsr与sr_TN和sr_TP的大小一致

import rasterio
from rasterio import warp  # 导入 warp 模块
from rasterio.enums import Resampling
from rasterio.warp import reproject, Resampling
import rasterio
from rasterio.enums import Resampling
import numpy as np





def resample_to_match(target_src, source_src):
    """将源栅格重采样至目标栅格的坐标系和尺寸"""
    destination = np.zeros(target_src.shape, dtype=np.float32)
    reproject(
        source=source_src.read(1),
        destination=destination,
        src_transform=source_src.transform,
        src_crs=source_src.crs,
        dst_transform=target_src.transform,
        dst_crs=target_src.crs,
        resampling=Resampling.nearest
    )
    return destination
def calculate_soil_conservation_value(Qsr_path, sr_TN_path, sr_TP_path, mud_siltation_coefficient=0.24, soil_bulk_density=1.27, reservoir_siltation_cost=60, TN_treatment_cost=3000, TP_treatment_cost=3500):
    """
    计算土壤保持价值。

    参数:
    Qsr_path (str): 土壤保持的实物量文件路径。
    sr_TN_path (str): 污染物氮纯含量文件路径。
    sr_TP_path (str): 污染物磷纯含量文件路径。
    mud_siltation_coefficient (float): 泥沙淤积系数。默认为0.24。
    soil_bulk_density (float): 土壤容重。默认为1.27。
    reservoir_siltation_cost (float): 水库单位清淤工程费用。默认为60。
    TN_treatment_cost (float): 污染物氮单位处理成本。默认为3000。
    TP_treatment_cost (float): 污染物磷单位处理成本。默认为3500。

    返回:
    Vsd (numpy.ndarray): 减少泥沙淤积价值。
    Vdpd (numpy.ndarray): 减少面源污染价值。
    Vsr (numpy.ndarray): 土壤保持价值。
    """

    # 读取土壤保持的实物量
    with rasterio.open(Qsr_path) as src_Qsr:
        Qsr = src_Qsr.read(1)
        qsr_meta = src_Qsr.meta

    # 读取污染物氮纯含量
    with rasterio.open(sr_TN_path) as src_sr_TN:
        sr_TN = resample_to_match(src_Qsr, src_sr_TN)

    with rasterio.open(sr_TP_path) as src_sr_TP:
        sr_TP = resample_to_match(src_Qsr, src_sr_TP)

    # 计算减少泥沙淤积价值
    Vsd = mud_siltation_coefficient * Qsr / soil_bulk_density * reservoir_siltation_cost

    # 计算减少面源污染价值
    Vdpd = Qsr * (sr_TN * TN_treatment_cost / 1000 + sr_TP * TP_treatment_cost / 1000)

    # 计算土壤保持价值
    Vsr = Vsd + Vdpd

    return Vsd, Vdpd, Vsr

def save_array_as_tiff(array, output_path, src):
    """
    将数组保存为TIF文件。

    参数:
    array (numpy.ndarray): 要保存的数组。
    output_path (str): 输出文件路径。
    src (rasterio.io.DatasetReader): 源文件，用于获取元数据。
    """
    with rasterio.open(
            output_path,
            'w',
            driver='GTiff',
            height=array.shape[0],
            width=array.shape[1],
            count=1,
            dtype=array.dtype,
            crs=src.crs,
            transform=src.transform,
    ) as dst:
        dst.write(array, 1)


# 示例用法
Qsr_path = r'E:\GEP_data\sr\Qsr\lingchuan_merged_Qsr.tif'
sr_TN_path =r'E:\GEP_data\sr\lingchuan_merged_tn_Pro.tif'
sr_TP_path = r'E:\GEP_data\sr\lingchuan_merged_tp_Pro.tif'

with rasterio.open(Qsr_path) as src:
    Vsd, Vdpd, Vsr = calculate_soil_conservation_value(Qsr_path, sr_TN_path, sr_TP_path)

    save_array_as_tiff(Vsd, 'Vsd.tif', src)
    save_array_as_tiff(Vdpd, 'Vdpd.tif', src)
    save_array_as_tiff(Vsr, 'Vsr.tif', src)